import numpy as np
import pandas as pd
import cv2
from tqdm import tqdm
from Preprocess import IMG_Path,LB_Path,CSV_Path,mkDir,YOLO_PATH
import shutil
from Preprocess import mkDir,getCsv,div_TrainData
import yolov5
from yolov5 import train
import os

EPOCHS = 20
IMG_SIZE=736
MODEL_DIR_NAME='model'
TRAIN_PIC_SIZE=640
BATCH_SIZE = 12

def trainModel():
    train.run(data='data.yaml', imgsz=TRAIN_PIC_SIZE, weights=YOLO_PATH+'yolov5s.pt',batch_size=BATCH_SIZE,epochs=EPOCHS,project=os.path.join(os.getcwd(),MODEL_DIR_NAME))


if __name__ == "__main__":
    # mkDir(IMG_Path, LB_Path)
    # df=getCsv(IMG_Path, LB_Path)
    # df.to_csv(CSV_Path, encoding="utf-8")
    # train_meta=div_TrainData(delete_src=True)
    trainModel()